﻿import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.version
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np

input_size = 28  # 图像的总尺寸28*28
numclasses = 10  # 标签的种类数
num_epochs = 3  # 训练的总循环周期
batch_size = 64  # 一个撮(批次)的大小，64张图片
learning_rate = 0.001  # 学习率

current_working_directory = os.getcwd()
data_file_path = os.path.join(current_working_directory, "Dataset", "Minist", "File")

train_dataset = datasets.MNIST(
    root=data_file_path,
    train=True,
    transform=transforms.ToTensor(),
    # , download=True
)

# 测试集
test_dataset = datasets.MNIST(
    root=data_file_path, train=False, transform=transforms.ToTensor()
)

# 构建batcb数据
train_loader = torch.utils.data.DataLoader(
    dataset=train_dataset, batch_size=batch_size, shuffle=True
)
test_loader = torch.utils.data.DataLoader(
    dataset=test_dataset, batch_size=batch_size, shuffle=True
)


class NN(nn.Module):
    def __init__(self):
        super(NN, self).__init__()
        # 定义第一层全连接层，将输入784映射到256
        self.fc1 = nn.Linear(784, 256)
        # 定义第二层全连接层，将256映射到10
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        # 首先将输入展平为一维向量
        x = x.view(x.size(0), -1)  # 假设输入是(batch_size, 1, 28, 28)
        # 通过第一层全连接层并应用ReLU激活函数
        x = torch.nn.functional.relu(self.fc1(x))
        # 通过第二层全连接层得到最终输出
        output = self.fc2(x)
        return output


# 示例：创建模型实例
model = NN()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):

        # image = images[0,:,:,:]
        # image_np = image.numpy().squeeze()

        # # 显示图像
        # plt.imshow(image_np, cmap='gray')
        # plt.title(f'Label: {labels[0].item()}')
        # plt.show()

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if (i + 1) % 100 == 0:
            print(
                f"Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(train_loader)}], Loss: {loss.item():.4f}"
            )

mechine_file_path = os.path.join(
    current_working_directory,
    "Pytorch",
    "Network",
    "Nn",
    "Machine",
    "model_parameters.pth",
)
torch.save(model.state_dict(), mechine_file_path)
# # 测试模型
# model.eval()  # 将模型设置为评估模式
# with torch.no_grad():
#     correct = 0
#     total = 0
#     for images, labels in test_loader:
#         outputs = model(images)
#         _, predicted = torch.max(outputs.data, 1)
#         total += labels.size(0)
#         correct += (predicted == labels).sum().item()
#     print(
#         f"Accuracy of the model on the 10000 test images: {100 * correct / total:.2f}%"
#     )